from pathlib import Path
from math import sqrt
import gmsh_api.gmsh as gmsh
import os
import sys
import re


def dist(p0, p1):
    return sqrt((p0.x - p1.x)**2 + (p0.y - p1.y)**2 + (p0.z - p1.z)**2)


class PhysicalGroup:
    def __init__(self):
        self.name = ""
        self.elem_list = []
        self.node_list = {}


class Element(object):
    def __init__(self):
        self.id = -1
        self.type = 0
        self.info = ""
        self.plist = []


class Node(object):
    def __init__(self):
        self.id = 0
        self.x = 0.
        self.y = 0.
        self.z = 0.
        self.info = ""
        self.dof = ""
        self.svalue = 0.
        self.evalue = 0.


class Mesh(PhysicalGroup):
    """docstring for mesh"""

    def __init__(self):
        PhysicalGroup.__init__(self)


class Crack(Mesh):
    """docstring for crack"""

    def __init__(self):
        Mesh.__init__(self)


class Boundary(Mesh):
    """docstring for Boundary"""

    def __init__(self):
        Mesh.__init__(self)


class Phase(object):
    def __init__(self):
        self.m_dirichlet = PhysicalGroup()
        self.m_neumann = PhysicalGroup()


class gmsh_parser(object):
    m_elem_types = {
        1: "segment",
        2: "triangle",
        3: "quadrangle",
        4: "tetrahedron",
        5: "hexahedron",
        6: "prism",
        7: "pyramid",
        8: "segment",
        9: "triangle",
        16: "quadrangle"
    }
    """Docstring for gmsh_parser. """

    def __init__(self):
        self.m_domain = PhysicalGroup()
        self.m_phases = {}
        self.m_cracks = {}
        self.m_mats = {}
        self.m_node_num = 0
        self.m_elem_num = 0

    def importGmsh(self, file_name_with_path):
        print("\n------------------------------------")
        if not Path(file_name_with_path).is_absolute():
            ScriptPath = os.path.split(os.path.realpath(sys.argv[0]))[0]
            os.chdir(ScriptPath)
            self.m_filename = Path(ScriptPath + "/" + file_name_with_path)
        else:
            self.m_filename = Path(file_name_with_path)
        if not self.m_filename.is_file():
            print("File \"", self.m_filename, "\" not exists! Program exited.")
            sys.exit()

        gmsh.initialize()
        print("0. parsing \"", self.m_filename, "\" .....\n")
        gmsh.open(self.m_filename.name)
        self.dim = gmsh.model.getDimension()
        gmsh.model.mesh.renumberNodes()
        gmsh.model.mesh.renumberElements()

        for p_dim, p_tag in gmsh.model.getPhysicalGroups():
            name = gmsh.model.getPhysicalName(p_dim, p_tag)
            if "mat" in name:
                self.addDomain(p_dim, p_tag)
            elif "crack" in name:
                self.addCrack(p_dim, p_tag)
            elif "dirichlet" in name:
                self.addDirichlet(p_dim, p_tag)
            elif "neumann" in name:
                self.addNeumann(p_dim, p_tag)

    def addDomain(self, p_dim, p_tag):
        p = self.parsePhysicalGroup(p_dim, p_tag)
        matchObj = re.match(".*mat(\\d+)_(\\w+).*", p.name)
        if not matchObj:
            raise Exception("domain object tag", p.name, "is invalid")
        mat_id = int(matchObj.group(1))
        mat_name = matchObj.group(2)
        self.m_mats[mat_id] = mat_name

        self.m_domain.elem_list += p.elem_list
        self.m_domain.node_list.update(p.node_list)

        print("  new material #", mat_id, mat_name, "added!")
        print("  ", len(p.elem_list), "new elements added!")
        print("  ", len(p.node_list), "new nodes added!")

    def addDirichlet(self, p_dim, p_tag):
        p = self.parsePhysicalGroup(p_dim, p_tag)
        # ------------- dirichlet boundary value --------------
        matchObj = re.match("dirichlet_P(\\d+)_.*_.*_(.*)_(.*)", p.name)
        if not matchObj:
            raise Exception("dirichlet object tag", p.name, "is invalid")
        phase_id = int(matchObj.group(1))
        svalue = float(matchObj.group(2))
        evalue = float(matchObj.group(3))
        for node in p.node_list.values():
            node.svalue = svalue
            node.evalue = evalue
        if phase_id not in self.m_phases:
            self.m_phases[phase_id] = Phase()
        self.m_phases[phase_id].m_dirichlet.node_list.update(p.node_list)

        print("  ", len(p.node_list), "new dirichlet node added to phase",
              phase_id, "!")

    def addNeumann(self, p_dim, p_tag):
        p = self.parsePhysicalGroup(p_dim, p_tag)
        # ------------- neumann boundary value --------------
        matchObj = re.match("neumann_P(\\d+)_.*_.*_(.*)_(.*)", p.name)
        if not matchObj:
            raise Exception("neumann object tag", p.name, "is invalid")
        phase_id = int(matchObj.group(1))
        svalue = float(matchObj.group(2))
        evalue = float(matchObj.group(3))
        self.m_phases[phase_id].m_neumann.node_list.update(p.node_list)
        for elem in p.elem_list.values():
            if elem.type == 1:  # linear segment neumann value
                p0 = self.m_phases[phase_id].m_neumann.node_list[elem.plist[0]]
                p1 = self.m_phases[phase_id].m_neumann.node_list[elem.plist[1]]
                length = dist(p0, p1)
                p0.svalue += length * 0.5 * svalue
                p0.evalue += length * 0.5 * evalue
                p1.svalue += length * 0.5 * svalue
                p1.evalue += length * 0.5 * evalue
            elif elem.type == 8:  # quad segment neumann value
                p0 = self.m_phases[phase_id].m_neumann.node_list[elem.plist[0]]
                p1 = self.m_phases[phase_id].m_neumann.node_list[elem.plist[1]]
                p2 = self.m_phases[phase_id].m_neumann.node_list[elem.plist[2]]
                length = dist(p0, p1)
                p0.svalue += length / 6. * svalue
                p0.evalue += length / 6. * evalue
                p1.svalue += length / 6. * svalue
                p1.evalue += length / 6. * evalue
                p2.svalue += length * .666666 * svalue
                p2.evalue += length * .666666 * evalue
            else:
                raise Exception("invalid element type", elem.type)
        if phase_id not in self.m_phases:
            self.m_phases[phase_id] = Phase()
        self.m_phases[phase_id].m_neumann.node_list.update(p.node_list)

        print("  ", len(p.node_list), "new neumann node added to phase",
              phase_id, "!")

    def addCrack(self, p_dim, p_tag):
        p = self.parsePhysicalGroup(p_dim, p_tag)
        matchObj = re.match(".*crack(\\d+)_.*", p.name)
        if not matchObj:
            raise Exception("domain object tag", p.name, "is invalid")
        crack_id = matchObj.group(1)
        if crack_id not in self.m_cracks:
            self.m_cracks[crack_id] = PhysicalGroup()
        self.m_cracks[crack_id].elem_list += p.elem_list
        self.m_cracks[crack_id].node_list.update(p.node_list)

    def write(self, path=""):
        print("\n1. writing mesh...")
        self.writeMesh()
        print("\n2. writing materials...")
        self.writeMaterials()
        print("\n3. writing boundary conditions...")
        self.writeBoundary()
        print("\n4. writing crack...")
        self.writeCrack()
        self.save()

    def parsePhysicalGroup(self, p_dim, p_tag):
        p = PhysicalGroup()
        p.name = gmsh.model.getPhysicalName(p_dim, p_tag)
        entity_tags = gmsh.model.getEntitiesForPhysicalGroup(p_dim, p_tag)
        for entity_tag in entity_tags:
            elem_types, elem_tags, node_tags_in_entity = gmsh.model.mesh.getElements(
                p_dim, entity_tag)
            type_id = 0
            # ------------- get elems in each entity -------------
            for elem_type in elem_types:
                (_, _, _, elem_nn, _,
                 _) = gmsh.model.mesh.getElementProperties(elem_type)
                for elem_id in range(len(elem_tags[type_id])):
                    elem = Element()
                    elem.type = elem_type
                    elem.info = p.name
                    for nn_id in range(elem_nn):
                        node_tag = node_tags_in_entity[type_id][
                            elem_id * elem_nn + nn_id] - 1
                        elem.plist.append(int(node_tag))
                    p.elem_list.append(elem)
                type_id += 1
            # ------------- get nodes in each entity -------------
            node_tags, coords, _ = gmsh.model.mesh.getNodes(
                p_dim, entity_tag, True)
            matchObj = re.match(".*_P\\d+_(.*?)_.*", p.name)
            for node_id in range(len(node_tags)):
                node = Node()
                node.id = int(node_tags[node_id] - 1)
                node.x = coords[node_id * 3 + 0]
                node.y = coords[node_id * 3 + 1]
                node.z = coords[node_id * 3 + 2]
                node.info = p.name
                if matchObj:
                    node.dof = matchObj.group(1)
                p.node_list[int(node.id)] = node
        return p

    def save(self, path_str=""):
        pass

    def writeMesh(self):
        pass

    def writeMaterials(self):
        pass

    def writeBoundary(self):
        pass

    def writeCrack(self):
        pass


if (__name__ == "__main__"):
    obj = gmsh_parser()
    obj.importGmsh("example.msh")
